#include <iostream>
#include <fstream>
#include <vector>
#include <string>

#include "particle/Particle.h"
#include "particle/PID.h"
#include "event/BelleEvent.h"
#include "tuple/BelleTupleManager.h"
#include "basf/module.h"
#include "basf/module_descr.h"
#include "benergy/BeamEnergy.h"
#include "kid/atc_pid.h"

#include "eid/eid.h"

#include "panther/panther.h"
#include MDST_H
#include EVTCLS_H
#include BELLETDF_H

#include "mdst/Muid_mdst.h"

#include "rows/ypipi_pi0row.hpp"
#include "rows/ypipi_ypi0pi0row.hpp"
#include <pytpp.hpp>

#include "ypipi_utils.h"

using std::string;
using std::vector;
using std::cout;
using std::endl;

#ifdef BELLE_NAMESPACE
namespace Belle {
#endif //BELLE_NAMESPACE

  class ypi0pi0_h5: public Module {
  public:
    ypi0pi0_h5();
    virtual ~ypi0pi0_h5();
    void init(int *);
    void term();
    void disp_stat(const char*);
    void hist_def();
    void event(BelleEvent*, int*);
    void begin_run(BelleEvent*, int*);
    void end_run(BelleEvent*, int*);

    char outpath_stub[4096];
    char inpath[4096]; 

  private:
    float calc_R2();

    int _nevents_tot;
    int _event_yield;
    int _trackno;

    std::ofstream _ofs_r2;

    pyt::table<ypipi_pi0row_t>* _nt_pi0;
    ypipi_pi0row_t* _row_pi0;
    pyt::table<ypipi_ypi0pi0row_t>* _nt;
    ypipi_ypi0pi0row_t* _row;

  };

  extern "C" Module_descr *mdcl_ypi0pi0_h5()
  {

    ypi0pi0_h5 *tmp = new ypi0pi0_h5();
    Module_descr *dscr = new Module_descr ( "ypi0pi0_h5", tmp);
    dscr->define_param("outpath_stub", "Path stub to output file", "S", 4096, tmp->outpath_stub);
    dscr->define_param("inpath", "Input file", "S", 4096, tmp->inpath);

    BeamEnergy::define_global(dscr);

    return dscr;
  }
  
  ypi0pi0_h5::ypi0pi0_h5() : 
  _event_yield(0), 
    _row_pi0(new ypipi_pi0row_t),
    _row(new ypipi_ypi0pi0row_t)
  {
    // Initialize member vars
  }

  ypi0pi0_h5::~ypi0pi0_h5() {
    // Delete member vars
    delete _nt_pi0;
    delete _row_pi0;
    delete _nt;
    delete _row;

  }

  void ypi0pi0_h5::init(int* status) {
    string path_stub(outpath_stub);

    string outpath = path_stub + ".h5";
    string pi0path = path_stub + ".pi0.h5";

    _nt = pyt::open_table<ypipi_ypi0pi0row_t>(outpath, "nt", "w");
    _nt->table_setattr("mdst_path", inpath);
    _nt->table_setattr("module", "ypipi");

    _nt_pi0 = pyt::open_table<ypipi_pi0row_t>(pi0path, "nt", "w");
    _nt_pi0->table_setattr("mdst_path", inpath);
    _nt_pi0->table_setattr("module", "ypipi.pi0");

    _ofs_r2.open("r2comp.txt");
  }
  
  void ypi0pi0_h5::hist_def() {
    // Define ntuples here
  }

  void ypi0pi0_h5::begin_run(BelleEvent* be, int* status) {
    BeamEnergy::begin_run();
  }
 
  void ypi0pi0_h5::event(BelleEvent *evptr, int *status) {
    Belle_event_Manager& m_event = Belle_event_Manager::get_manager();

    _row->expno   = m_event[0].ExpNo();
    _row->runno   = m_event[0].RunNo();
    _row->eventno = m_event[0].EvtNo();
    _row->beam_e  = BeamEnergy::E_beam_corr();

    vector<Particle> v_mup;
    vector<Particle> v_mum;
    vector<Particle> v_pi0;
  
    Ptype ptype_mup("MU+");
    Ptype ptype_mum("MU-");

    atc_pid pid_pi_vs_K(3,1,5,3,2);
    atc_pid pid_mu_vs_K(3,1,5,1,3);
    Mdst_charged_Manager& m_charged = Mdst_charged_Manager::get_manager();
    for (Mdst_charged_Manager::iterator it = m_charged.begin();
	 it != m_charged.end(); ++it)
      {
	Mdst_trk& trk(it->trk());
	Muid_mdst muid(*it);

	if (muid.Muon_likelihood() < 0.1)
	  continue;
	
	if (it->charge() > 0.0) {
	  Particle tmp(*it, ptype_mup);
	  v_mup.push_back(tmp);
	  _row->mup_muid_llh = muid.Muon_likelihood();
	  _row->mup_pion_vs_kaon_llh = pid_pi_vs_K.prob(*it);
	  _row->mup_muon_vs_kaon_llh = pid_mu_vs_K.prob(*it);
	}
	else {
	  Particle tmp(*it, ptype_mum);
	  v_mum.push_back(tmp);
	  _row->mum_muid_llh = muid.Muon_likelihood();
	  _row->mum_pion_vs_kaon_llh = pid_pi_vs_K.prob(*it);
	  _row->mum_muon_vs_kaon_llh = pid_mu_vs_K.prob(*it);
	}
    }

    if ((v_mup.size() != 1) || (v_mum.size() != 1))
      return;

    Mdst_pi0_Manager& m_pi0 = Mdst_pi0_Manager::get_manager();
    Mdst_ecl_aux_Manager& m_ecl_aux = Mdst_ecl_aux_Manager::get_manager();

    _row->pi0_start_index = _trackno;
    Hep3Vector cm_boost = BeamEnergy::CMBoost();

    for (Mdst_pi0_Manager::iterator it = m_pi0.begin();
	 it != m_pi0.end(); it++)
      {
	
	if (fabs(0.135 - it->mass()) >= 0.0054*5) 
	  continue;
		
	Mdst_ecl& ecl_g0 = it->gamma(0).ecl();
	Mdst_ecl& ecl_g1 = it->gamma(1).ecl();
		
	// if (ecl_g0.energy() < 0.1)
	//   continue;

	// if (ecl_g1.energy() < 0.1)
	//   continue;

 	Mdst_ecl_aux& ecl_aux_g0 = m_ecl_aux(ecl_g0.get_ID());
	Mdst_ecl_aux& ecl_aux_g1 = m_ecl_aux(ecl_g1.get_ID());
	
	// if (ecl_aux_g0.e9oe25() < 0.95)
	//   continue;
	
	// if (ecl_aux_g1.e9oe25() < 0.95)
	//   continue;

	_row_pi0->chisq    = it->chisq();
	_row_pi0->energy   = it->energy();
	_row_pi0->g0e9oe25 = ecl_aux_g0.e9oe25();
	_row_pi0->g1e9oe25 = ecl_aux_g1.e9oe25();
	_row_pi0->g0energy = ecl_g0.energy();
	_row_pi0->g1energy = ecl_g1.energy();
	_row_pi0->mass     = it->mass();
	_row_pi0->px       = it->px();
	_row_pi0->py       = it->py();
	_row_pi0->pz       = it->pz();
	_row_pi0->trackno  = _trackno;

	_row_pi0->eventno = m_event[0].EvtNo();
	
	_trackno++;

	_nt_pi0->append(*_row_pi0);
	
	Particle tmp(*it);
	v_pi0.push_back(tmp);
      }

    _row->npi0 = _trackno - _row->pi0_start_index;

    if (v_pi0.size() < 2)
      return;

    HepLorentzVector Y5S = HepLorentzVector(0, 0, 0, 10.876);
    
    HepLorentzVector pi01, pi02;
    if (v_pi0[0].e() > v_pi0[1].e()) {      
      pi01 = v_pi0[0].p();
      pi02 = v_pi0[1].p();
    }
    else {
      pi01 = v_pi0[1].p();
      pi02 = v_pi0[0].p();
      Particle tmp = v_pi0[1];
      v_pi0[1] = v_pi0[0];
      v_pi0[0] = tmp;
    }
      
    HepLorentzVector mup = v_mup[0].p();
    HepLorentzVector mum = v_mum[0].p();

    pi01.boost(-cm_boost);
    pi02.boost(-cm_boost);
    mup.boost(-cm_boost);
    mum.boost(-cm_boost);

    _row->mup_cm_px = mup.px();
    _row->mup_cm_py = mup.py();
    _row->mup_cm_pz = mup.pz();
    _row->mup_cm_e  = mup.e();

    _row->mum_cm_px = mum.px();
    _row->mum_cm_py = mum.py();
    _row->mum_cm_pz = mum.pz();
    _row->mum_cm_e  = mum.e();

    _row->pi01_cm_px = pi01.px();
    _row->pi01_cm_py = pi01.py();
    _row->pi01_cm_pz = pi01.pz();
    _row->pi01_cm_e  = pi01.e();

    _row->pi02_cm_px = pi02.px();
    _row->pi02_cm_py = pi02.py();
    _row->pi02_cm_pz = pi02.pz();
    _row->pi02_cm_e  = pi02.e();

    Mdst_pi0 mdst_pi01 = v_pi0[0].mdstPi0();
    
    _row->pi01_m     = mdst_pi01.mass();
    _row->pi01_chisq = mdst_pi01.chisq();
    
    Mdst_ecl& pi01_ecl1 = mdst_pi01.gamma(0).ecl();
    Mdst_ecl& pi01_ecl2 = mdst_pi01.gamma(1).ecl();
    
    if (pi01_ecl1.energy() < pi01_ecl2.energy()) {
      Mdst_ecl& tmp = pi01_ecl1;
      pi01_ecl1 = pi01_ecl2;
      pi01_ecl2 = tmp;
    }

    _row->pi01_g1e = pi01_ecl1.energy();  
    _row->pi01_g2e = pi01_ecl2.energy();    

    Mdst_ecl_aux& pi01_ecl_aux1 = m_ecl_aux(pi01_ecl1.get_ID());
    Mdst_ecl_aux& pi01_ecl_aux2 = m_ecl_aux(pi01_ecl2.get_ID());
    
    _row->pi01_g1e9oe25 = pi01_ecl_aux1.e9oe25();
    _row->pi01_g2e9oe25 = pi01_ecl_aux2.e9oe25();

    Mdst_pi0 mdst_pi02 = v_pi0[1].mdstPi0();
    
    _row->pi02_m     = mdst_pi02.mass();
    _row->pi02_chisq = mdst_pi02.chisq();
    
    Mdst_ecl& pi02_ecl1 = mdst_pi02.gamma(0).ecl();
    Mdst_ecl& pi02_ecl2 = mdst_pi02.gamma(1).ecl();
    
    _row->pi02_g1e = pi02_ecl1.energy();
    _row->pi02_g2e = pi02_ecl2.energy();

    Mdst_ecl_aux& pi02_ecl_aux1 = m_ecl_aux(pi02_ecl1.get_ID());
    Mdst_ecl_aux& pi02_ecl_aux2 = m_ecl_aux(pi02_ecl2.get_ID());
    
    _row->pi02_g1e9oe25 = pi02_ecl_aux1.e9oe25();
    _row->pi02_g2e9oe25 = pi02_ecl_aux2.e9oe25();

    _row->r2 = calc_R2();
    
    HepLorentzVector initial(0, 0, 0, 10.876);

    _row->mm_off_pi0pi0 = (initial - pi01 -pi02).m();
    _row->m_mumu        = (mup + mum).m();

    Evtcls_hadron_info_Manager& m_hadron_info = Evtcls_hadron_info_Manager::get_manager();

    if (m_hadron_info.count() > 0)
      _ofs_r2 << _row->eventno << "\t" << _row->r2 << "\t" << m_hadron_info[0].R2() << endl;

    _nt->append(*_row);

    _event_yield++;
        
  }

  float ypi0pi0_h5::calc_R2() {
    Mdst_charged_Manager& m_charged = Mdst_charged_Manager::get_manager();
    Mdst_gamma_Manager& m_gamma = Mdst_gamma_Manager::get_manager();
    
    
    int n_trk = m_charged.count() + m_gamma.count();
    float H0, R1, R2, R3, R4;
    float p_trk[n_trk*3];

    float beam_e = BeamEnergy::E_beam_corr();

    int trkno = 0;
    for (Mdst_charged_Manager::iterator it = m_charged.begin(); 
	 it != m_charged.end(); it++, trkno++)
      {
	p_trk[3*trkno]     = it->px();
	p_trk[3*trkno + 1] = it->py();
	p_trk[3*trkno + 2] = it->pz();
      }

    for (Mdst_gamma_Manager::iterator it = m_gamma.begin();
	 it != m_gamma.end(); it++, trkno++)
      {
	p_trk[3*trkno]     = it->px();
	p_trk[3*trkno + 1] = it->py();
	p_trk[3*trkno + 2] = it->pz();
      }	

    fwjet2(n_trk, p_trk, beam_e, &H0, &R1, &R2, &R3, &R4);

    return R2; 
  }


  void ypi0pi0_h5::end_run(BelleEvent* be, int* status) { 
    // Yessir
  }

  void ypi0pi0_h5::term() {
    _ofs_r2.close();
    cout << "CRAIG : " << _event_yield << " mu+mu-pi0pi0 events found" << endl;
    // End of analysis, close ntuples
  }

  void ypi0pi0_h5::disp_stat(const char* str) {
    // Who knows
  }

#ifdef BELLE_NAMESPACE
}
#endif
